package com.zhang.wanx.common.util;

import com.zhang.wanx.common.model.GenerationReqDTO;
import com.zhang.wanx.common.model.GenerationResDTO;
import com.zhang.wanx.common.model.InputDTO;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import javax.annotation.PostConstruct;

/**
 * description
 *
 * @author zhangFanJun
 * @date 2024-01-29 11:10
 **/
@Slf4j
@Data
@RequiredArgsConstructor
@Component
public class WangXiangRpc {

    private WebClient webClient;

    @PostConstruct
    public void init() {
        webClient = WebClient.builder()
                .baseUrl("https://dashscope.aliyuncs.com/api/v1")
                .build();
    }


    /**
     * styleIndex：
     * 0 复古漫画
     * 1 3D童话
     * 2 二次元
     * 3 小清新
     * 4 未来科技
     * 5 国画古风
     * 6 将军百战
     * 7 炫彩卡通
     * 8 清雅国风
     * 9 喜迎新年
     *
     * @param imageUrl 图片地址
	 * @param styleIndex 类型
     * @return 结果
     * @author zfj
     * @date 2024/3/27
     */
    public GenerationResDTO generation(String imageUrl, Integer styleIndex) {

        GenerationReqDTO in = new GenerationReqDTO();
        in.setModel("wanx-style-repaint-v1");
        InputDTO inputDTO = new InputDTO();
        inputDTO.setStyle_index(styleIndex);
        inputDTO.setImage_url(imageUrl);
        in.setInput(inputDTO);

        GenerationResDTO result = webClient.post()
                .uri("/services/aigc/image-generation/generation")
                .bodyValue(in)
                .headers(h -> {
                    h.add("Content-Type", "application/json");
                    h.add("X-DashScope-Async", "enable");
                    h.add("Authorization", "Bearer sk-3fa79c45a7f9467781c387f5c7c5c5b9");
                })
                .accept(MediaType.APPLICATION_JSON)
                .exchangeToMono(response -> {
                    HttpStatus status = response.statusCode();
                    if (status.is2xxSuccessful()) {
                        return response.bodyToMono(GenerationResDTO.class);
                    } else {
                        return response.createException().flatMap(Mono::error);
                    }
                }).block();

        log.info("返回结果：",result);

        return result;
    }


    public GenerationResDTO tasks(String id) {

        GenerationResDTO result = webClient.get()
                .uri("/tasks/"+id)
                .headers(h -> {
                    h.add("Content-Type", "application/json");
                    h.add("Authorization", "Bearer sk-3fa79c45a7f9467781c387f5c7c5c5b9");
                })
                .accept(MediaType.APPLICATION_JSON)
                .exchangeToMono(response -> {
                    HttpStatus status = response.statusCode();
                    if (status.is2xxSuccessful()) {
                        return response.bodyToMono(GenerationResDTO.class);
                    } else {
                        return response.createException().flatMap(Mono::error);
                    }
                }).block();

        log.info("返回结果：",result);

        return result;
    }
}
